// "Copyright [2021] <Copyright Zhancheng Liang, Shengkai Lin>"

// The following CPU module function as a controller of a node
// It is responsible for receiving messages from the central controller
// and forward the route info to the App and router

#include <fstream>
#include <iomanip>
#include <string>
#include "cpu.h"

namespace ai_sim {

Define_Module(Cpu);

void Cpu::initialize() {
    my_global_index = par("my_global_index");
    my_index = par("my_index");
    need_wait_topo = par("need_wait_topo");
    omnetpp::cModule * gpu_module =
        getParentModule()->getSubmodule("gpu", my_index - 1);
    my_gpu = dynamic_cast<Gpu *> (gpu_module);
    secure_time = par("secure_time");
    assert(my_gpu->getIndex() == getIndex() - 1);
}

void Cpu::handleReuqires(Flow& the_flow) {
    // Handle change route require
    if (the_flow.requires[NeedControllerCooperate]) {
        EV << "handle controller cooperate" << std::endl;
        // the_flow.requires[NeedControllerCooperate] = 0;
        int clustertype = ClusterType(getSystemModule()->par("clustertype").intValue());

        if (!the_flow.dependRoute.empty() && 
            (clustertype == ClusterType::OpticalCluster || clustertype == ClusterType::HybridCluster)) {
            FlowMsg *f_msg = new FlowMsg;
            f_msg->setThe_flow(the_flow);
            f_msg->setKind(CHANGE_ROUTE_MSG);
            sendDelayed(f_msg, secure_time, "controller_port$o");
            // send(f_msg, "controller_port$o");
        } else {
            SwitchAckMsg* sa_msg = new SwitchAckMsg;
            sa_msg->setThe_flow(the_flow);
            scheduleAt(omnetpp::simTime(), sa_msg);
        }
    }
    // Handle compute require
    // it = the_flow.requires.find(NeedCompute);
    if (the_flow.requires[NeedCompute]) {
        EV << "handle compute" << std::endl;
        // the_flow.requires[NeedCompute] = false;
        FlowMsg *f_msg = new FlowMsg;
        f_msg->setThe_flow(the_flow);
        scheduleAt(omnetpp::simTime() + the_flow.the_task.compute_time,
            f_msg);
    }
}

// Check if the requirements has all satisfied
// and then send them out oro  handle the requirements
void Cpu::alloc_flow(int tx_flow_id) {
    // The communication between cpu and gpu is reduced to direct function
    // call
    if (running_task.tx_flows.find(tx_flow_id) != running_task.tx_flows.end()) {
        Flow& tx_flow = running_task.tx_flows[tx_flow_id];
        int remainRequire = 0;
        for (auto require : tx_flow.requires) {
            remainRequire += require.second;
        }
        if (tx_flow.numUpstreamflows != 0)
            error("The flow still have upstreams!");
        if (remainRequire == 0) {
             startTime[std::make_pair(tx_flow.srcIdx , tx_flow.destIdx )] = omnetpp::simTime();
            std::ofstream ofile;
            int numMachine = getParentModule()->getParentModule()->par("numMachine");
            int modelsize_ = getParentModule()->getParentModule()->par("modelsize_");
            ofile.open(std::to_string(numMachine)+"starthd_ECN"+ std::to_string(modelsize_) +".txt",ios::app);
            ofile << tx_flow.srcIdx + 1 << "-" << tx_flow.destIdx + 1  << "-" << omnetpp::simTime() <<endl;
            std::cout<<"flow from "<< tx_flow.srcIdx<<" to "<< tx_flow.destIdx<<" with datasize "<<tx_flow.data_size<<" start at "<<startTime[std::make_pair(tx_flow.srcIdx , tx_flow.destIdx )] <<endl;
            ofile.close();
            my_gpu->handleSend(tx_flow);
            running_task.tx_flows.erase(tx_flow_id);
        } else {
            handleReuqires(tx_flow);
        }
    }
}

void Cpu::handleReceive(const FlowMsg *msg) {
    Enter_Method("handleReceive(FlowMsg *msg)");
    Flow r_flow = msg->getThe_flow();
    // Erase the flow in the rx_flows map
    if (running_task.rx_flows.find(r_flow.flowId) !=
        running_task.rx_flows.end()) {
        running_task.rx_flows.erase(r_flow.flowId);
        // std::cout << running_task.rx_flows.size() << endl;
    } else {
        std::cout << "Rec flow to " << r_flow.destIdx <<
            " on " << my_gpu->par("my_global_index").intValue() << endl;
        std::cout << "The src is " << r_flow.srcIdx << endl;
        std::cout << "the flow id is " << r_flow.flowId << endl;
        error("ERROR MSG RECEIVED");
    }
    std::ofstream ofile;
    int numMachine = getParentModule()->getParentModule()->par("numMachine");
    int modelsize_ = getParentModule()->getParentModule()->par("modelsize_");
    ofile.open(std::to_string(numMachine)+"resulthd_ECN"+ std::to_string(modelsize_) +".txt",ios::app);
    ofile << r_flow.srcIdx + 1 << "-" << r_flow.destIdx + 1  << "-" << omnetpp::simTime() <<endl;
    //std::cout<<"flow from "<< r_flow.srcIdx<<" to "<< r_flow.destIdx<<" with datasize "<<r_flow.data_size<<" finish at "<< omnetpp::simTime()  <<endl;
    ofile.close();
    if (running_task.rx_flows.empty()) {
        handleFinish(msg);
        EV << "handleFinish on cpu" << my_index << std::endl;
    }
    std::vector<int> d_flows_id = r_flow.downstream_flows;
    if (d_flows_id.empty()) return;
    for (int d_flow_id : d_flows_id) {
        try {
            Flow& downstreamflow = running_task.tx_flows.at(d_flow_id);
            if (--downstreamflow.numUpstreamflows <= 0) {
                // Send out the downstream flow
                alloc_flow(d_flow_id);
            }
        } catch (std::out_of_range) {
            error("No such flow in the CPU, id %d", d_flow_id);
            // EV_WARN << "No such flow in the CPU" << std::endl;
        }
    }
}

void Cpu::handleFinish(const FlowMsg *msg) {
    Flow r_flow = msg->getThe_flow();
    // TaskMsg *f_msg = new TaskMsg;
    // f_msg->setNow_task(r_flow.the_task);
    EV << "flow from " << r_flow.srcIdx<<" to "<< r_flow.destIdx<<" with datasize "<<r_flow.data_size<<" finish at "<< omnetpp::simTime()  <<endl;
    FlowMsg *finish_msg = new FlowMsg;
    finish_msg->setThe_flow(msg->getThe_flow());
    finish_msg->setKind(FINISHED_MSG);
    send(finish_msg, "controller_port$o");
}

void Cpu::handleNewTask(const TaskFlowVecMsg *msg) {
    ASSERT(msg->getTargetId() == getId());
    running_task.now_task = msg->getNow_task();
    // Put the tx_flows in the unordered_map
    for (auto& flow : msg->getRx_flows()) {
        running_task.rx_flows[flow.flowId] = flow;
    }

    // Put the tx_flows in the unordered_map
    for (auto& flow : msg->getTx_flows()) {
        // EV << "numUpstreamflows" << flow.numUpstreamflows << endl;

        if (flow.outportId == -1) {
            error("Stop1");
        }
        running_task.tx_flows[flow.flowId] = flow;
        if (flow.numUpstreamflows == 0) {
        // If the flow has no upstreamflows, alloc it.
            alloc_flow(flow.flowId);
        }
    }

    EV << "Recieved a TASK_ALLOC_MSG at cpu"
        << my_global_index << "\t The num flows is "
        << msg->getTx_flows().size() << std::endl;
}

void Cpu::handleMessage(omnetpp::cMessage *msg) {
    /////////////////// If the cpu is a normal cpu////////////////
    if (my_index != 0) {
        if (msg->getKind() == SWITCH_ACK_MSG) {
            SwitchAckMsg * oxc_ack_msg = dynamic_cast<SwitchAckMsg *> (msg);
            int the_flow_id = oxc_ack_msg->getThe_flow().flowId;
            running_task.tx_flows[the_flow_id].
                requires[NeedControllerCooperate] = 0;
            alloc_flow(the_flow_id);
            delete oxc_ack_msg;
        } else if (msg->getKind() == TASK_ALLOC_MSG) {
            // A new task arrived
            TaskFlowVecMsg *t_msg = dynamic_cast<TaskFlowVecMsg *>(msg);
            handleNewTask(t_msg);
            delete t_msg;
        } else if (msg->isSelfMessage()) {
            // Between batches, finish compute
            FlowMsg *f_msg = dynamic_cast<FlowMsg *>(msg);
            SwitchAckMsg * sa_msg = dynamic_cast<SwitchAckMsg*>(msg);
            Flow* the_flow;
            int the_flow_id;
            if (f_msg != nullptr) {
                the_flow_id = f_msg->getThe_flow().flowId;
                running_task.tx_flows[the_flow_id].
                    requires[NeedCompute] = 0;
                the_flow = &running_task.tx_flows[the_flow_id];
            } else if (sa_msg != nullptr) {
                the_flow_id = sa_msg->getThe_flow().flowId;
                running_task.tx_flows[the_flow_id].
                    requires[NeedControllerCooperate] = 0;
                the_flow = &running_task.tx_flows[the_flow_id];
            } else {
                error("unknown self kind");
            }
            int remainRequire = 0;
            for (auto require : the_flow->requires) {
                remainRequire += require.second;
            }
            if (remainRequire == 0) alloc_flow(the_flow_id);
            delete msg;
        } else {
            int i = msg->getKind();
            delete msg;
            error("DONT SUPPORT THE MSG TYPE %d ON NORMAL CPU", i);
        }
    } else {
        //////////////// Else this is a master cpu.///////////////
        switch (msg->getKind()) {
            case SWITCH_ACK_MSG: {
                SwitchAckMsg * oxc_ack_msg = dynamic_cast<SwitchAckMsg *> (msg);
                int srcIdx = oxc_ack_msg->getThe_flow().srcIdx;
                send(oxc_ack_msg, "cpu_port$o",
                    srcIdx % gateSize("cpu_port$o"));
                return;
            }
            case TASK_ALLOC_MSG: {
                TaskFlowVecMsg *t_msg = dynamic_cast<TaskFlowVecMsg *>(msg);
                // Get the srcIdx from the flow
                int target;
                if (!t_msg->getTx_flows().empty())
                    target = t_msg->getTx_flows()[0].srcIdx;
                else if (!t_msg->getRx_flows().empty())
                    target = t_msg->getRx_flows()[0].destIdx;
                else
                    error("Empty tx_flows and rx_flows");
                send(t_msg, "cpu_port$o", target % gateSize("cpu_port$o"));
                break;
            }
            case FINISHED_MSG : {
                send(msg, "controller_port$o");
                break;
            }
            case CHANGE_ROUTE_MSG : {
                send(msg, "controller_port$o");
                break;
            }
            default: {
                delete msg;
                error("DONOT SUPPORT THE MSG TYPE %d ON MASTER CPU",
                    msg->getKind());
                break;
            }
        }
    }
}
}  // namespace ai_sim
